#include "core/memory.h"
#include "tools/log.h"
#include "comm/boot_info.h"
#include "tools/klib.h"
#include "cpu/mmu.h"
#include "dev/console.h"

static bitmap_manager_t bitmap_manager_up1MB;
static page_directory_entry_t kernel_page_dir[PAGE_DIRECTORY_LENGTH] __attribute__((aligned(MEMORY_PAGE_SIZE)));

static void bitmap_manager_init(bitmap_manager_t * bitmap_manager, 
                    uint32_t start, 
                    uint32_t size, 
                    uint32_t page_size,
                    uint8_t * bits)
{
    mutex_init(&bitmap_manager->mutex);
    bitmap_manager->start = start;
    bitmap_manager->page_size = page_size;
    bitmap_manager->size = size;
    bitmap_init(&bitmap_manager->bitmap, bits, size / page_size, 0);
}

// 从位图管理器中分配page_count页 返回分配出的首页的物理地址
static uint32_t bitmap_manager_alloc_page(bitmap_manager_t * bitmap_manager, int page_count)
{
    uint32_t addr = 0;
    mutex_lock(&bitmap_manager->mutex);
    int page_index = bitmap_alloc_nbits(&bitmap_manager->bitmap, 0, page_count);
    if(page_index >= 0 )
    {
        addr = bitmap_manager->start + page_index * bitmap_manager->page_size;
    }
    mutex_unlock(&bitmap_manager->mutex);
    return addr;
}

// 从位图管理器中释放page_count页
static void bitmap_manager_free_page(bitmap_manager_t * bitmap_manager, uint32_t physcial_addr, int page_count)
{
    mutex_lock(&bitmap_manager->mutex);
    int page_index = (physcial_addr - bitmap_manager->start) / bitmap_manager->page_size;
    bitmap_bit_set(&bitmap_manager->bitmap, page_index, 0, page_count);
    mutex_unlock(&bitmap_manager->mutex);
}

void show_memory_info(boot_info_t * boot_info)
{
    log_printf("RAM REGION:");
    for(int i=0;i < boot_info->ram_region_count; i++)
    {
        log_printf("region %d:", i);
        log_printf("    start %x:", boot_info->ram_region_cfg[i].start);
        log_printf("    size %x:", boot_info->ram_region_cfg[i].size);
    }
}

static uint32_t total_memory_size(boot_info_t * boot_info)
{
    uint32_t memory_size = 0;
    for(int i=0;i < boot_info->ram_region_count; i++)
    {
        memory_size += boot_info->ram_region_cfg[i].size;
    }
    return memory_size;
}

// 创建用户使用的线性地址、目前用于创建进程的页表
uint32_t memory_create_user_linear_mem(void)
{
    // 这里分配的是用户页表 是一级页表
    page_directory_entry_t * page_dir = (page_directory_entry_t *)bitmap_manager_alloc_page(&bitmap_manager_up1MB,1);
    if(!page_dir)
    {
        return 0;
    }
    
    kernel_memset(page_dir,0,MEMORY_PAGE_SIZE);
    
    uint32_t user_pde_start = index_in_pde(MEMORY_TASK_BASE);
    // 用户页表 前部分和操作系统一样
    for(int i = 0; i < user_pde_start; i++)
    {
        page_dir[i].v = kernel_page_dir[i].v;
    }

    return (uint32_t)page_dir;
}

// 从页目录表中找寻虚拟地址对应的二级页表项，is_alloc表示未找到页表是否分配
 page_table_entry_t * find_page_table(page_directory_entry_t * kernel_page_dir,
                                    uint32_t linear_addr,
                                    int is_alloc)
 {
    page_table_entry_t * page_table = 0;
    // 获取线性地址对应的一级页表表项
    page_directory_entry_t * pde = kernel_page_dir + index_in_pde(linear_addr);

    if(pde->persent)
    {
        // 获取表项对应的二级页表地址
        uint32_t pte_addr = (pde->physical_page_table_addr << 12);  
        page_table = (page_table_entry_t *)pte_addr;
    }
    else if(is_alloc)
    {
        // 如果没有找到二级页表，则新分配一个二级页表，返回该页表起始地址
        // 这些地址是对齐4kb的，低12位没用
        uint32_t pte_addr = bitmap_manager_alloc_page(&bitmap_manager_up1MB,1);
        // 页表分配失败
        if(!pte_addr)
        {
            return 0;
        } 
        // 分配成功，将二级页表起始地址 写入页目录表项
        pde->v = pte_addr | PDE_P | PDE_W | PDE_U;
        page_table = (page_table_entry_t *) pte_addr;
        kernel_memset(page_table, 0, MEMORY_PAGE_SIZE);
    }
    // 从线性地址中间10位获取 该线性地址在二级页表的地址返回
    return page_table + index_in_pte(linear_addr);
 }
 
// 将线性地址start->线性地址end，映射到物理页start
int memory_create_map_virtual_to_physcial(page_directory_entry_t * kernel_page_dir, 
                            uint32_t linear_addr,
                            uint32_t physical_addr,
                            int page_count,
                            int permission)
{
    for(int i = 0; i < page_count; i++)
    {
        // log_printf("create map: v-%x p-%x, permission: %x", linear_addr, physical_addr, permission);
        // 从一级页表中找物理页对应的二级页表项
        page_table_entry_t * pte = find_page_table(kernel_page_dir, linear_addr, 1);
        if(!pte)
        {
            log_printf("create pte failed. pte == 0");
            return -1;
        }

        // log_printf("\tpte addr: %x", (uint32_t)pte);
        ASSERT(pte->persent == 0);
        pte->v = physical_addr | permission | PTE_P;
        
        linear_addr += MEMORY_PAGE_SIZE;
        physical_addr += MEMORY_PAGE_SIZE;
    }
}

// 内核这块内存 放入页表
void create_kernel_table(void)
{
    // 从lds文件中获取各区起始结束地址
    extern uint8_t text_start[],rodata_end[],data_start[];
    extern uint8_t kernel_start[];
    static memory_map_t memory_map_table[5] = {
        // 线性起始地址，线性结束地址，物理起始地址，是否可写
        // 前三行是1M一下的位置 ，到0x80000结束，逻辑地址和物理地址完全一致
        {kernel_start, text_start, kernel_start, PTE_W},
        {text_start, rodata_end, text_start,PTE_R},
        {data_start, (void *)MEMORY_EBDA_START, data_start,PTE_W},
        // 控制台所控制的显存地址，从0xb8000开始，之后32kb的区域
        {(void *)CONSOLE_DISP_ADDR, (void *)CONSOLE_DISP_END, (void *)CONSOLE_DISP_ADDR, PTE_W},
        // 1M以上的空间
        {(void *)MEMORY_EXT_START, (void *)MEMORY_EXT_END, (void *)MEMORY_EXT_START, PTE_W},

    };

    for(int i=0;i < sizeof(memory_map_table) / sizeof(memory_map_t);i++)
    {
        memory_map_t * memory_map = memory_map_table + i;
        uint32_t linear_addr_start = align_bound_down((uint32_t)memory_map->linear_addr_start, 
                                                        MEMORY_PAGE_SIZE);

        uint32_t linear_addr_end = align_bound_up((uint32_t)memory_map->linear_addr_end,
                                                    MEMORY_PAGE_SIZE);

        int page_count = (linear_addr_end - linear_addr_start) / MEMORY_PAGE_SIZE;
    
        // 将线性地址start->线性地址end，映射到物理页start
        memory_create_map_virtual_to_physcial(kernel_page_dir,
                                linear_addr_start, 
                                (uint32_t)memory_map->physical_addr_start, 
                                page_count,
                                memory_map->permission);
    }
}

void memory_init(boot_info_t * boot_info)
{

    log_printf("memory init......");
    show_memory_info(boot_info);
    
    // 1M以上 空闲的内存
    uint32_t memory_up1MB_free_size = total_memory_size(boot_info) - MEMORY_EXT_START;
    // 一页是4kb,需要把空闲内存对齐页面大小，好分配
    memory_up1MB_free_size = align_bound_down(memory_up1MB_free_size,MEMORY_PAGE_SIZE); 

    log_printf("free memory size: %x", memory_up1MB_free_size);

    // lds文件中定义的地址,用于保存位图
    extern uint8_t * mem_free_start;
    uint8_t * mem_free = (uint8_t *)&mem_free_start;

    bitmap_manager_init(&bitmap_manager_up1MB, 
                        MEMORY_EXT_START, 
                        memory_up1MB_free_size, 
                        MEMORY_PAGE_SIZE, 
                        mem_free);

    // 位图指针跳过位图 指向下一个空闲内存区
    mem_free += bitmap_byte_count(&bitmap_manager_up1MB.bitmap);
    ASSERT(mem_free < (uint8_t *)MEMORY_EBDA_START);
    
    // 内核内存放入页表
    create_kernel_table();
    mmu_set_page_dir((uint32_t)kernel_page_dir);
}

// 为指定页表分配空间
int memory_alloc_page_for_page_dir(
        uint32_t page_dir, 
        uint32_t virtual_addr, // 线性地址
        uint32_t alloc_size, 
        int permission)
{
    int page_count = align_bound_up(alloc_size, MEMORY_PAGE_SIZE) / MEMORY_PAGE_SIZE;
    for(int i = 0; i < page_count; i++)
    {
        uint32_t physcial_addr = bitmap_manager_alloc_page(&bitmap_manager_up1MB,1);
        if(!physcial_addr)
        {
            log_printf("memory alloc page for task's page dir failed");
            return -1;
        }

        int err_code = memory_create_map_virtual_to_physcial(
            (page_directory_entry_t *)page_dir,
            virtual_addr,
            physcial_addr,
            1,
            permission);
        if(err_code < 0)
        {
            log_printf("memory create map virtual to physcial failed! error code : %d", err_code);
            return -1;
        }
        virtual_addr += MEMORY_PAGE_SIZE;
    }

    return 0;
}

// 为进程分配物理页面
int memory_alloc_page_for(uint32_t start_addr, uint32_t alloc_size, int permission)
{
    return memory_alloc_page_for_page_dir( 
                curr_task_get()->tss.cr3, 
                start_addr, 
                alloc_size, 
                permission
            );
}

uint32_t memory_alloc_page(void)
{
    uint32_t physcial_addr = bitmap_manager_alloc_page(&bitmap_manager_up1MB,1);
    if(!physcial_addr)
    {
        log_printf("memory alloc page for task's kernel stack failed");
        return -1;
    }

    return physcial_addr;
}

static page_directory_entry_t * curr_page_dir(void)
{
    return (page_directory_entry_t *)(curr_task_get()->tss.cr3);
}

void memory_free_page(uint32_t page_addr)
{
    if(page_addr < MEMORY_TASK_BASE)
    {
        bitmap_manager_free_page(&bitmap_manager_up1MB,page_addr,1);
    }
    else
    {
        page_table_entry_t * pte = find_page_table(curr_page_dir(),page_addr,0);
        ASSERT(pte && (pte->persent));
        bitmap_manager_free_page(&bitmap_manager_up1MB,(uint32_t)pte,1);
        pte->v = 0;
    }
}

void memory_free_user_linear_mem(uint32_t page_dir)
{
    uint32_t user_pde_start = index_in_pde(MEMORY_TASK_BASE);
    page_directory_entry_t * pde = (page_directory_entry_t *)page_dir + user_pde_start;
    ASSERT(page_dir != 0);

    for(int i=user_pde_start;i<PAGE_DIRECTORY_LENGTH;i++,pde++)
    {
        if(!pde->persent) continue;
        page_table_entry_t *pte = (page_table_entry_t *)pte_addr_in_pde(pde);
        for(int j=0;j<PAGE_TABLE_LENGTH;j++,pte++)
        {
            if(!pte->persent) continue;
            bitmap_manager_free_page(&bitmap_manager_up1MB, page_addr_in_pte(pte),1);
        }
        bitmap_manager_free_page(&bitmap_manager_up1MB, (uint32_t)pte, 1);

    }
    bitmap_manager_free_page(&bitmap_manager_up1MB, (uint32_t)page_dir, 1);
}

uint32_t memory_copy_user_linear_mem(uint32_t page_dir)
{
    uint32_t new_page_dir = memory_create_user_linear_mem();
    if(!new_page_dir)
    {
        goto memory_copy_user_linear_mem_failed; 
    }
    uint32_t user_pde_start = index_in_pde(MEMORY_TASK_BASE);
    page_directory_entry_t * pde = (page_directory_entry_t *)page_dir + user_pde_start;
    
    for(int i=user_pde_start;i<PAGE_DIRECTORY_LENGTH;i++,pde++)
    {
        if(!pde->persent) continue;

        page_table_entry_t *pte = (page_table_entry_t *)pte_addr_in_pde(pde);

        for(int j=0;j<PAGE_TABLE_LENGTH;j++,pte++)
        {
            if(!pte->persent) continue;

            uint32_t new_page_table = bitmap_manager_alloc_page(&bitmap_manager_up1MB,1);
            
            if(!new_page_table)
            {
                goto memory_copy_user_linear_mem_failed;
            } 
            
            uint32_t vitual_addr = (i<<22) | (j<<12);   
            int err = memory_create_map_virtual_to_physcial((page_directory_entry_t *)new_page_dir,
                                                            vitual_addr,
                                                            new_page_table,
                                                            1,
                                                            get_pte_perm(pte));
            if(err < 0)
            {
                goto memory_copy_user_linear_mem_failed;
            }
            kernel_memcpy((void *)new_page_table, (void *)vitual_addr, MEMORY_PAGE_SIZE);
        }
    }
    return new_page_dir;

memory_copy_user_linear_mem_failed:
    if(new_page_dir){
        memory_free_user_linear_mem(new_page_dir);
    } 
    return 0;
}

uint32_t memory_get_pyhscial_addr_from_page_dir(uint32_t page_dir, uint32_t vaddr)
{
    page_table_entry_t * pte = find_page_table((page_directory_entry_t *)page_dir, vaddr, 0);
    if(!pte)
    {
        return 0;
    }
    //   页面起始地址 + 偏移
    return page_addr_in_pte(pte) + (vaddr & (MEMORY_PAGE_SIZE - 1));
}

int memory_copy_user_linear_mem_data(uint32_t aim_addr, 
                                    uint32_t new_page_dir, 
                                    uint32_t from_addr, 
                                    int size)
{
    while (size > 0)
    {
        uint32_t aim_physcial_addr = memory_get_pyhscial_addr_from_page_dir(new_page_dir,aim_addr);
        if(!aim_physcial_addr)
        {
            return -1;
        }

        int offset_in_page = aim_physcial_addr & (MEMORY_PAGE_SIZE - 1);
        int copy_size = MEMORY_PAGE_SIZE - offset_in_page;
        if(copy_size > size)
        {
            copy_size = size;
        }
        kernel_memcpy((void *)aim_physcial_addr, (void *)from_addr, copy_size);
        size -= copy_size;
        from_addr += copy_size;
        aim_addr += copy_size;
    }
}  

char * sys_sbrk(int incr)
{
    ASSERT(incr >= 0);
    int pre_incr = incr;
    task_t * task = curr_task_get();
    char * pre_heap_end = (char *)task->heap_end;
    if(incr == 0)
    {
        log_printf("sbrk(0): end = %x", pre_heap_end);
        return (char *)pre_heap_end;
    }

    uint32_t start = task->heap_end;
    uint32_t end = start + incr;

    uint32_t offset = start % MEMORY_PAGE_SIZE;
    if(offset)
    {
        if(offset + incr <= MEMORY_PAGE_SIZE)
        {
            task->heap_end = end;
            // log_printf("sbrk(%d): end = %x", pre_incr, end);
            return pre_heap_end;
        }
        else
        {
            uint32_t curr_size = MEMORY_PAGE_SIZE - offset;
            start += curr_size;
            incr -= curr_size;
        }
    }

    if(incr)
    {
        uint32_t curr_size = end - start;
        int err = memory_alloc_page_for(start, curr_size, PTE_P | PTE_U | PTE_W); 
        if(err < 0)
        {
            log_printf("sbrk alloc memory failed!");
            return (char *)-1;
        }
    }
    
    // log_printf("sbrk(%d): end = %x", pre_incr, end);
    task->heap_end = end;
    return pre_heap_end;   
}